import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8mb4'
}

engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@"
    f"{db_config['host']}:{db_config['port']}/{db_config['database']}?"
    f"charset={db_config['charset']}"
)

# 1. 先查询表中所有股票代码和日期范围，确认数据是否存在
print("正在检查表中是否有数据...")
check_df = pd.read_sql_query("""
    SELECT DISTINCT ts_code, MIN(trade_date) as min_date, MAX(trade_date) as max_date 
    FROM date_1 
    WHERE trade_date BETWEEN '2023-01-01' AND '2023-12-31'
    LIMIT 10  -- 显示前10个股票代码
""", engine)

if check_df.empty:
    print("错误：2023年全年在date_1表中无任何数据！")
else:
    print("表中存在的股票代码及日期范围：")
    print(check_df)

# 2. 根据检查结果修改股票代码（例如华夏银行正确代码是600015.SH）
target_ts_code = '600015.SH'  # 修正：将SZ改为SH（银行股多数在沪市）
print(f"\n尝试查询股票代码：{target_ts_code}")

df = pd.read_sql_query(f"""
    SELECT * FROM date_1 
    WHERE trade_date BETWEEN '2023-01-01' AND '2023-12-31' 
      AND ts_code = '{target_ts_code}'
""", engine)

if df.empty:
    print(f"仍然无数据！请从上述列表中选择存在的ts_code替换")
else:
    print(f"成功读取 {len(df)} 行数据")
    # 计算涨跌幅（确保closes是数值型）
    df['closes'] = pd.to_numeric(df['closes'], errors='coerce')
    df = df.dropna(subset=['closes'])
    df['rd_closer'] = round((df['closes'] - df['closes'].shift(1)) / df['closes'].shift(1), 2)
    df = df.dropna(subset=['rd_closer'])
    
    # 回归分析
    ex = ['id', 'ts_code', 'trade_date', 'the_date']
    number_list = [col for col in df.select_dtypes(include=['number']).columns if col not in ex]
    if number_list:
        formula = 'rd_closer ~ ' + ' + '.join(number_list)
        res = smf.ols(formula=formula, data=df).fit()
        print("\n回归结果：")
        print(res.summary())
